package com.crux.ribbon;

import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Gauge;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

public class CruxLoadBalancerLifecycle implements LoadBalancerLifecycle<Object, Object, ServiceInstance> {
    private final MeterRegistry meterRegistry;
    private final ConcurrentHashMap<ServiceInstance, AtomicLong> activeRequestsPerInstance = new ConcurrentHashMap();
    //每个实例的响应时间
    private final ConcurrentHashMap<ServiceInstance, Long> responseTimesPerInstance = new ConcurrentHashMap();

    public ConcurrentHashMap<ServiceInstance, AtomicLong> getActiveRequestsPerInstanceUsed() {
        return this.activeRequestsPerInstance;
    }

    public ConcurrentHashMap<ServiceInstance, Long> getResponseTimesPerInstance() {
        return responseTimesPerInstance;
    }

    public long getResponseTimeByInstanceId(String instanceId) {
        //return responseTimesPerInstance;
        return responseTimesPerInstance.keySet().stream()
                .filter(instance -> instance.getInstanceId().equalsIgnoreCase(instanceId))
                .map(responseTimesPerInstance::get).findFirst().orElseGet(() -> 0L);
    }

    private ConcurrentHashMap<ServiceInstance, AtomicLong> activeRequestsPerInstanceUsed = new ConcurrentHashMap<>();

    public CruxLoadBalancerLifecycle(MeterRegistry meterRegistry) {
        this.meterRegistry = meterRegistry;
    }

    @Override
    public void onStart(Request<Object> request) {

    }

    @Override
    public void onStartRequest(Request<Object> request, Response<ServiceInstance> lbResponse) {
        if (request.getContext() instanceof TimedRequestContext) {
            ((TimedRequestContext) request.getContext()).setRequestStartTime(System.nanoTime());
        }

        if (lbResponse.hasServer()) {
            ServiceInstance serviceInstance = (ServiceInstance) lbResponse.getServer();
            AtomicLong activeRequestsCounter = (AtomicLong) this.activeRequestsPerInstance.computeIfAbsent(serviceInstance, (instance) -> {
                AtomicLong createdCounter = new AtomicLong();
                Gauge.builder("loadbalancer.requests.active", () -> {
                    return createdCounter;
                }).tags(MyLoadBalancerTags.buildServiceInstanceTags(serviceInstance)).register(this.meterRegistry);
                return createdCounter;
            });
            activeRequestsCounter.incrementAndGet();
        }
    }

    @Override
    public void onComplete(CompletionContext<Object, ServiceInstance, Object> completionContext) {
        long requestFinishedTimestamp = System.nanoTime();
        if (CompletionContext.Status.DISCARD.equals(completionContext.status())) {
            Counter.builder("loadbalancer.requests.discard").tags(MyLoadBalancerTags.buildDiscardedRequestTags(completionContext)).register(this.meterRegistry).increment();
        } else {
            ServiceInstance serviceInstance = (ServiceInstance) completionContext.getLoadBalancerResponse().getServer();
            AtomicLong activeRequestsCounter = (AtomicLong) this.activeRequestsPerInstance.get(serviceInstance);
//            if (activeRequestsCounter != null) {
//                activeRequestsCounter.decrementAndGet();
//            }

            Object loadBalancerRequestContext = completionContext.getLoadBalancerRequest().getContext();
            if (this.requestHasBeenTimed(loadBalancerRequestContext)) {
                long responseTime = requestFinishedTimestamp - ((TimedRequestContext) loadBalancerRequestContext).getRequestStartTime();
                responseTimesPerInstance.put(serviceInstance, requestFinishedTimestamp);
                if (CompletionContext.Status.FAILED.equals(completionContext.status())) {
                    Timer.builder("loadbalancer.requests.failed").tags(MyLoadBalancerTags.buildFailedRequestTags(completionContext)).register(this.meterRegistry).record(responseTime, TimeUnit.NANOSECONDS);
                    return;
                }
                Timer.builder("loadbalancer.requests.success").tags(MyLoadBalancerTags.buildSuccessRequestTags(completionContext)).register(this.meterRegistry).record(responseTime, TimeUnit.NANOSECONDS);
            }

        }
    }

    private boolean requestHasBeenTimed(Object loadBalancerRequestContext) {
        return loadBalancerRequestContext instanceof TimedRequestContext && ((TimedRequestContext) loadBalancerRequestContext).getRequestStartTime() != 0L;
    }
}
